from core.decorators.database_decorator import with_async_session
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, asc,delete,update

from core.dependencies import AsyncSessionLocal
from core.utils.paginator import Paginator
from src.modules.knowbase.repository.entity.vectors import TextVector
from src.modules.knowbase.schemas.model_schema import TextVectorResponse
from typing import List

class TextVectorDao:
    @staticmethod
    @with_async_session()
    async def get_textvector_list(user_id:str = None, file_id:str = None,
                                   kb_id:str = None,
                                   page: int = 1, page_size:int = 10,
                                   session:AsyncSession = None):
        query = select(TextVector)
        if user_id:
            query = query.where(TextVector.user_id == user_id)
        if file_id:
            query = query.where(TextVector.file_id == file_id)
        if kb_id:
            query = query.where(TextVector.know_base_id == kb_id)
        query = query.order_by(asc(TextVector.uid))
        paginator  = Paginator(session = session,
                               model = TextVector,
                               response_model = TextVectorResponse,
                               base_query=query)
        return await paginator.paginate(page=page, page_size = page_size)



    @staticmethod
    @with_async_session()
    async def delete_by_file_id(user_id:str = None,
                                file_id:str = None,
                                session: AsyncSession = None):
        if not file_id :
            raise ValueError("参数校验失败,缺少file_id！")
        delete_sql = delete(TextVector).where(TextVector.file_id == file_id)
        result =  await session.execute(delete_sql)
        return result.rowcount>0



    @staticmethod
    @with_async_session()
    async def delete_by_kb_id(user_id:str = None,
                              kb_id:str = None,
                              session: AsyncSession = None):
        if not kb_id:
            raise ValueError("参数校验失败，缺少kb_id")
        delete_sql = delete(TextVector).where(TextVector.know_base_id == kb_id)
        result = await session.execute(delete_sql)
        return result.rowcount>0



    @staticmethod
    async def get_textvector_list_for_rag(question_vector: List[float],
                                          session: AsyncSession,
                                          knowbase_ids: List[str],
                                          top_k: int):
        if not knowbase_ids:
            return []
        query = (
            select(
                TextVector,
                TextVector.text_vector.l2_distance(question_vector).label('distance')
            )
            .where(TextVector.know_base_id.in_(knowbase_ids))
            .order_by('distance')
            .limit(top_k)
        )
        result = await session.execute(query)
        text_results = result.all()
        return text_results


    @staticmethod
    async def delete_vector(vector_id:str):
        session:AsyncSession = AsyncSessionLocal()
        try:
            delete_sql = delete(TextVector).where(TextVector.vector_id == vector_id)
            result = await session.execute(delete_sql)
            await session.commit()
            return result.rowcount>0
        except Exception as e:
            await session.rollback()
            return False
        finally:
            await session.close()

    @staticmethod
    async def update_text_vector_by_id(
            vector_id:str,
            new_text_content:str,
            new_text_vector:list
    ):
        session: AsyncSession = AsyncSessionLocal()
        try:
            # 查询要更新的记录
            stmt = select(TextVector).where(TextVector.vector_id == vector_id)
            result = await session.execute(stmt)
            text_vector_obj = result.scalar_one_or_none()

            if not text_vector_obj:
                print(f"未找到vector_id为{vector_id}的记录")
                return False

            # 更新字段
            text_vector_obj.text_content = new_text_content
            text_vector_obj.text_vector = new_text_vector

            #提交更改
            await session.commit()
            print(f"成功更新vector_id为{vector_id}的记录")
            return True
        except Exception as e:
            await session.rollback()
            print(f"更新失败:{str(e)}")
            return False
        finally:
            await session.close()

            


